//
// Created by 12992 on 24-10-13.
//

#include "video_meta.h"
#include "codec_id_to_string.h"
#include "color_space_to_string.h"
#include <glog/logging.h>

extern "C"
{
#include <libswscale/swscale.h>
#include <libavformat/avformat.h>
#include <libavcodec/avcodec.h>
}

#include <filesystem>
#include <array>

namespace utils
{
    using namespace std;

    video_meta::video_meta(const std::string& input)
        : input_path(input)
    {
        try
        {
            // 检查路径合法性
            // if (!filesystem::is_regular_file(input))
            // {
            //     throw error::video_meta_error(error::video_meta_error::type::PATH_ERROR, "can not open file");
            // }
            this->video_name = filesystem::path(input).filename().string();
            //路径合法，打开文件
            int ret = avformat_open_input(&this->m_fmt_ctx, input.c_str(), nullptr, nullptr);
            if (ret < 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::FILE_OPEN_ERROR, ret);
            }

            //解析流信息
            ret = avformat_find_stream_info(this->m_fmt_ctx, nullptr);
            if (ret < 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::FIND_STREAM_ERROR, ret);
            }

            this->m_frames.resize(this->m_fmt_ctx->nb_streams, 0);

            // 初始化pkt
            this->m_pkt = av_packet_alloc();
            if (!this->m_pkt)
            {
                throw error::video_meta_error(error::video_meta_error::type::PACKET_ALLOC_ERROR,
                                              "can not alloc packet");
            }
            this->count_frame();
        }
        catch (error::video_meta_error& error)
        {
            this->clean_frame_counter(error.get_type());
            throw;
        }
        catch (...)
        {
            this->clean_frame_counter(error::video_meta_error::type::UNKNOWN_ERROR);
            throw;
        }
    }

    video_meta::~video_meta()
    {
        this->clean_frame_counter(error::video_meta_error::type::UNKNOWN_ERROR);
    }

    int video_meta::get_total_frame()
    {
        if (this->total_frame == -1)
        {
            this->count_frame();
        }
        return this->total_frame;
    }

    int video_meta::get_stream_frame(int index)
    {
        if (this->total_frame == -1)
        {
            this->count_frame();
        }
        return this->m_frames.at(index);
    }

    std::string video_meta::get_video_format() const
    {
        return this->m_fmt_ctx->iformat->name;
    }

    std::string video_meta::get_video_format2() const
    {
        return this->m_fmt_ctx->iformat->long_name;
    }

    std::string video_meta::get_video_encoder_type() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        const AVStream* video_stream = this->m_fmt_ctx->streams[index];
        return codec_id_to_string(video_stream->codecpar->codec_id);
    }

    int video_meta::get_video_framerate() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        auto framerate = this->m_fmt_ctx->streams[index]->avg_frame_rate;
        return framerate.num / framerate.den;
    }

    long long int video_meta::get_video_avg_bitrate() const
    {
        if (this->m_fmt_ctx->bit_rate > 0)
        {
            return this->m_fmt_ctx->bit_rate;
        }
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        const AVStream* video_stream = this->m_fmt_ctx->streams[index];
        LOG(INFO) << "bit rate : " << video_stream->codecpar->bit_rate;
        return video_stream->codecpar->bit_rate;
    }

    std::string video_meta::get_audio_encoder_type() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_AUDIO, -1, -1, nullptr, 0);
        const AVStream* audio_stream = this->m_fmt_ctx->streams[index];
        return codec_id_to_string(audio_stream->codecpar->codec_id);
    }

    std::string video_meta::get_audio_channel_layout() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_AUDIO, -1, -1, nullptr, 0);
        const AVStream* audio_stream = this->m_fmt_ctx->streams[index];
        const auto& layout = audio_stream->codecpar->ch_layout;
        switch (layout.nb_channels)
        {
            case 1: return "mono";
            case 2: return "stereo";
            case 3: return "2.1";
            default: return "unknown";
        }
    }

    int video_meta::get_audio_sample_rate() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_AUDIO, -1, -1, nullptr, 0);
        const AVStream* audio_stream = this->m_fmt_ctx->streams[index];
        return audio_stream->codecpar->sample_rate;
    }

    int video_meta::get_audio_stream_number() const
    {
        int res = 0;
        for (int i = 0; i < this->m_fmt_ctx->nb_streams; i++)
        {
            auto* curr_stream = this->m_fmt_ctx->streams[i];
            if (curr_stream->codecpar->codec_type == AVMEDIA_TYPE_AUDIO)
            {
                res++;
            }
        }
        return res;
    }

    std::string video_meta::get_filename() const
    {
        return this->video_name;
    }

    std::string video_meta::get_path() const
    {
        return this->input_path;
    }

    void video_meta::clean_frame_counter(error::video_meta_error::type _type)
    {
        switch (_type)
        {
            case error::video_meta_error::type::UNKNOWN_ERROR: av_packet_free(&this->m_pkt);
            case error::video_meta_error::type::PACKET_ALLOC_ERROR: ;
            case error::video_meta_error::type::FIND_STREAM_ERROR: avformat_close_input(&this->m_fmt_ctx);
            case error::video_meta_error::type::FILE_OPEN_ERROR: ;
            case error::video_meta_error::type::PATH_ERROR: ;
            default: ;
        }
    }

    void video_meta::count_frame()
    {
        bool flag = true;
        for (int i = 0; i < this->m_fmt_ctx->nb_streams; i++)
        {
            auto* curr_stream = this->m_fmt_ctx->streams[i];
            if (curr_stream->nb_frames > 0)
            {
                this->total_frame += curr_stream->nb_frames;
                this->m_frames[i] = curr_stream->nb_frames;
            }
            else
            {
                this->total_frame = 0;
                fill(this->m_frames.begin(), this->m_frames.end(), 0);
                flag = false;
                break;
            }
        }
        while (!flag)
        {
            int ret = av_read_frame(this->m_fmt_ctx, this->m_pkt);
            if (AVERROR_EOF == ret)
            {
                break;
            }
            if (ret < 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::READ_FRAME_ERROR, ret);
            }
            this->total_frame++;
            this->m_frames[this->m_pkt->stream_index]++;
            av_packet_unref(this->m_pkt);
        }
    }

    std::pair<int, int> video_meta::get_video_resolution() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        const AVStream* video_stream = this->m_fmt_ctx->streams[index];
        return {video_stream->codecpar->width, video_stream->codecpar->height};
    }

    long long int video_meta::get_duration() const
    {
        if (this->m_fmt_ctx->duration > 0)
        {
            return m_fmt_ctx->duration / AV_TIME_BASE;
        }
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        const AVStream* video_stream = this->m_fmt_ctx->streams[index];
        LOG(INFO) << "duration : " << video_stream->duration / video_stream->time_base.den;
        return video_stream->duration / video_stream->time_base.den;
    }

    std::string video_meta::get_video_color_space() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        const AVStream* video_stream = this->m_fmt_ctx->streams[index];
        auto fmt = static_cast<AVPixelFormat>(video_stream->codecpar->format);
        return color_space_to_string(fmt);
    }

    std::vector<utils::media_frame> video_meta::get_video_preview_frame() const
    {
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);

        // 初始化解码器
        const AVCodec* codec = avcodec_find_decoder(this->m_fmt_ctx->streams[index]->codecpar->codec_id);
        if (!codec)
        {
            throw error::video_meta_error(error::video_meta_error::type::DECODER_NOT_FOUND, "can not found decoder");
        }

        AVCodecContext* codec_ctx = avcodec_alloc_context3(codec);
        if (!codec_ctx)
        {
            throw error::video_meta_error(error::video_meta_error::type::CODER_CTX_ALLOC_ERROR,
                                          "can not alloc decoder ctx");
        }

        int ret = avcodec_parameters_to_context(codec_ctx, this->m_fmt_ctx->streams[index]->codecpar);
        if (ret < 0)
        {
            throw error::video_meta_error(error::video_meta_error::type::CODER_CTX_SET_ERROR, ret);
        }

        ret = avcodec_open2(codec_ctx, codec, nullptr);
        if (ret < 0)
        {
            throw error::video_meta_error(error::video_meta_error::type::CODER_OPEN_ERROR, ret);
        }
        vector<media_packet> packets;
        auto start_time = this->m_fmt_ctx->start_time;
        long long end_time;
        if (this->m_fmt_ctx->duration > 0)
        {
            end_time = this->m_fmt_ctx->duration + start_time;
            end_time = end_time / AV_TIME_BASE * this->m_fmt_ctx->streams[index]->time_base.den;
        }
        else
        {
            end_time = this->m_fmt_ctx->streams[index]->duration + start_time;
        }

        long long mid_time = (start_time + end_time) / 2;
        array<long long, 3> time_points{start_time, mid_time, end_time};
        for (long long tp: time_points)
        {
            ret = av_seek_frame(this->m_fmt_ctx, index, tp, AVSEEK_FLAG_BACKWARD);
            if (ret < 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::SEEK_FRAME_ERROR, ret);
            }
            packets.push_back(this->read_one_frame(index));
        }
        return this->flush_all_frame(codec_ctx, packets);
    }

    media_packet video_meta::read_one_frame(int index) const
    {
        int ret = 0;
        while (true)
        {
            ret = av_read_frame(this->m_fmt_ctx, this->m_pkt);
            if (ret < 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::READ_FRAME_ERROR, ret);
            }
            if (this->m_pkt->stream_index == index)
            {
                media_packet packet(this->m_pkt);
                LOG(INFO) << this->m_pkt->pts;
                av_packet_unref(this->m_pkt);
                return packet;
            }
            av_packet_unref(this->m_pkt);
        }
    }

    std::vector<utils::media_frame> video_meta::change_color_space(std::vector<utils::media_frame>& frames) const
    {
        // 初始化视频缩放对象
        const int index = av_find_best_stream(this->m_fmt_ctx, AVMEDIA_TYPE_VIDEO, -1, -1, nullptr, 0);
        auto [width, height] = this->get_video_resolution();
        SwsContext* ctx = sws_getContext(
            width,
            height,
            static_cast<AVPixelFormat>(this->m_fmt_ctx->streams[index]->codecpar->format),
            width, height, AV_PIX_FMT_RGB32,
            SWS_FAST_BILINEAR, nullptr, nullptr, nullptr);
        std::vector<utils::media_frame> vec;
        for (auto& frame: frames)
        {
            auto input = frame.get_frame();
            auto temp_frame = av_frame_alloc();
            temp_frame->width = width;
            temp_frame->height = height;
            temp_frame->format = AV_PIX_FMT_RGB32;
            int ret = av_frame_get_buffer(temp_frame, 0);
            if (ret < 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::SCALE_ERROR, ret);
            }
            ret = sws_scale(ctx, input->data, input->linesize, 0, input->height,
                            temp_frame->data, temp_frame->linesize);
            if (ret <= 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::SCALE_ERROR, ret);
            }
            vec.emplace_back(temp_frame);
            av_frame_free(&temp_frame);
        }
        sws_freeContext(ctx);
        return vec;
    }

    std::vector<utils::media_frame>
    video_meta::flush_all_frame(AVCodecContext* codec_ctx, std::vector<media_packet>& packets) const
    {
        packets.emplace_back();
        std::vector<utils::media_frame> vec;

        for (auto& packet: packets)
        {
            int ret = avcodec_send_packet(codec_ctx,
                                          bool(packet) ? packet.get_packet() : nullptr);
            if (ret < 0)
            {
                throw error::video_meta_error(error::video_meta_error::type::DECODE_ERROR, ret);
            }
            while (true)
            {
                auto frame = av_frame_alloc();
                ret = avcodec_receive_frame(codec_ctx, frame);
                if (AVERROR_EOF == ret || AVERROR(EAGAIN) == ret)
                {
                    av_frame_free(&frame);
                    break;
                }
                else if (ret < 0)
                {
                    throw error::video_meta_error(error::video_meta_error::type::DECODE_ERROR, ret);
                }
                vec.emplace_back(frame);
                av_frame_free(&frame);
            }
        }
        return this->change_color_space(vec);
    }
} // utils
